Install Acceleration Framework into Training Script - #157

from DM -

It would be great if there was a way to pass a Python dataclass instead of yaml file
python sft_trainer.py \ --acceleration_framework_config_file framework.yaml

we can create a new GPTQ-LOra config here https://github.com/foundation-model-stack/fms-hf-tuning/blob/main/tuning/config/peft_config.py#L21
so users of fms-tuning have exactly same experience as Lora

Inside fms-tuning, we know for GPTQ-Lora we need to call AccelerationFramework() -> can we convert to a Python dataclass object accepted by AccelerationFramework -AccelerationFramework(qLORAcONFIG )

You know what dataclass instance it is so you can parse params from that

we can implement this in steps though instead of requiring too much change to your framework-
Step 1 - users of fms-tuning continue to use dataclasses -> we will have to internally convert params passed to the yaml file you want (we wont expose yaml to users of tuning, but keep it inside tuning library ) -> call accelerationframework same way with yaml
step 2- > acceleration framework adds support for Python Dataclasses for each plugin as well besides yaml , so we dont have to create yaml but can just pass params needed for a certain plugin (edited)

if we do step 1, we can merge code and get it working, while working on step 2 as enhancement

We are discussing if we need 1 dataclass or multiple, but we can start with step 1 with dataclass only in tuning and see how it goes

@Ssukriti @alex-jw-brooks I have made the requested changes:

• now with dataclass parsing, it should be easy for the user to figure out what the required arguments are from inspecting the dataclasses. There is no more specifying of YAML.
• see for example quantized_lora_config.
• I have put unit tests in tests. It is almost complete but I might add one or two more.

Notes:

• it was not that straightforward to get transformers.HfArgumentParser work with nested dataclasses. I had to implement two additional utilities here. I can try to further simply this implementation.
  • made some attempts to simplify this with a decorator parsable_dataclass
• The argument parsing works as expected, you can see the tests

@fabianlim thank you for the redesign and the accelerationframework unit tests. Design looks good, so no major changes needed. Just few comments:

1. question above on when is GPTQLora applicable

2. The tests for acceleration framework check integration in detail, thank you. But I think if possible it would be beneficial to add some top level tests to test_sft_trainer as well, just to ensure
   a. if quantization_lora config passed, tuning still succeeds and model after tuning can still be loaded and inferred on (like rest of tuning unit tests) . I understand you may need a quantized base model to do these unit tests? so let me know if its feasible to add or not .
   b. similarly with kernels as well, if passed, tuning still succeeds.
   Either way might be good to add tests for failure case as well in test_sft_trainer. What if user passes GPTQLoraconfig to an unsupported model that is not quantized - what happens then? is error given and caught

3. Documentation may need some more updates - but we can do that in subsequent PRs. Mainly if there is any limitation on what model types we can apply QLoRA to (needs a 4bit quantized model), that should be documented and highlighted in README

@Ssukriti thank you for reviewing!

since this requires augmentation and peft_config , is quantized_lora_config only expected to work with LoRA tuning and Loraconfig , or can one also apply it with fine tuning and prompt tuning?

The AccelerationFramework has logic inside its plugins to check if the peft_config is not properly set. The peft_config is passed to the framework via the augmentation step. Maybe I can have a unit test to demonstrate this.

The tests for acceleration framework check integration in detail, thank you. But I think if possible it would be beneficial to add some top level tests to test_sft_trainer as well, just to ensure
a. if quantization_lora config passed, tuning still succeeds and model after tuning can still be loaded and inferred on (like rest of tuning unit tests) . I understand you may need a quantized base model to do these unit tests? so let me know if its feasible to add or not .
b. similarly with kernels as well, if passed, tuning still succeeds.
Either way might be good to add tests for failure case as well in test_sft_trainer. What if user passes GPTQLoraconfig to an unsupported model that is not quantized - what happens then? is error given and caught

Actually I think you are referring to this kind of test. We do already have one.

• I can add two much such tests, i) for fused ops and kernels, and ii) a negative test on an unsupported (non-quantized) model.

Documentation may need some more updates - but we can do that in subsequent PRs. Mainly if there is any limitation on what model types we can apply QLoRA to (needs a 4bit quantized model), that should be documented and highlighted in README

ok np!

@Ssukriti I have added a set of extra tests

These tests involve calling sft_trainer.train, so I think these are the tests that you are requesting for. They can be considered integration tests, to ensure correct working of the framework integrated with the trainer.

Therefore, I have added the following integration tests:

1. test_framework_raises_due_to_invalid_arguments: this test will demonstrate that the framework plugins will also check the arguments, and throw if invalid arguments are passed in. For example,. if we attempt to use accelerated peft and no peft_config is passed, then it throws
2. test_framework_intialized_properly_peft: this was refactored from an older test, but now it also demonstrates BNBQLora properly loading. Previously it was only GPTQ_LoRA happy path
3. test_framework_intialized_properly_foak: this one demonstrates that fused ops and kernels are also integrated properly.

@Ssukriti i have rebased the changes, and also created an issue here to track the remaining work items #205. Working on making the tests pass Update: all checks passing now.